from input import Input

from plugin_system import Plugin

import numpy as np

class XYZInput(Input):
    def __init__(self, filename=None, frame=0):
        self.filename = filename
        self.fl = None
        self.frame = frame
        
    def initialize(self):
        self.fl = open(self.filename)

    def finalize(self):
        self.fl.close()

    def read(self, system):
        for i in xrange(self.frame + 1):
            num_objects, step = map(int, next(self.fl).split())
            for j in xrange(num_objects):
                cols = next(self.fl).split()
                name = cols[0]
                positions = np.array(map(float, cols[1:4]), np.float64)
                velocities = np.array(map(float, cols[4:7]), np.float64)
                system.positions[name] = positions
                system.velocities[name] = velocities                         

class XYZInputPluginAdapter(Plugin, XYZInput):
    plugin_name = "XYZInput"
    parameter_definitions = {"filename" : (str, True, "File from which to read data."),
                             "frame" : (int, True, "Frame from which to read data. 0-indexed.")}
    
    def __init__(self, *args, **user_params):
        params = self.configure_parameters(**user_params)
        XYZInput.__init__(self, filename=params["filename"], frame=params["frame"])
        
